import asyncio
import time
import uuid
from datetime import datetime, timedelta, time
import logging
import os
import json
import websockets
import ControlSql as msql
import basicFuncation
from server_FastAPI import UPLOAD_FOLDER


class MicrosecondFormatter(logging.Formatter):
    def formatTime(self, record, datefmt=None):
        ct = self.converter(record.created)  # 将时间戳转换为时间元组
        if datefmt:
            if "%f" in datefmt:
                # 使用 datetime 支持微秒
                dt = datetime.fromtimestamp(record.created)
                return dt.strftime(datefmt)
            else:
                # 使用 time.strftime 格式化
                return time.strftime(datefmt, ct)
        else:
            # 默认格式
            t = time.strftime("%Y-%m-%d %H:%M:%S", ct)
            s = "%s,%03d" % (t, record.msecs)
            return s


# 配置 logging
formatter = MicrosecondFormatter(
    fmt="[%(asctime)s] %(levelname)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S,%f"  # 包含微秒
)
handler = logging.StreamHandler()
handler.setFormatter(formatter)

logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(handler)

# 存储WebSocket连接和文件信息
connected_clients = {}  # {client_id: websocket}
client_info = {}  # 存储客户端的事件
SERVER_IP = "127.0.0.1"
SERVER_FILES_DIR = os.path.join('E:\\project\\python\\DataStream\\Upload', "server_files")


async def check_client_online(client_id, websocket, disconnect_event):
    try:
        while True:
            await asyncio.sleep(30)  # 每隔30秒检查一次
            try:
                await websocket.ping()  # 发送ping消息
                # logging.info(f"Client {client_id} is still online.")
            except websockets.exceptions.ConnectionClosed:
                logging.info(f"Client {client_id} is offline.")
                disconnect_event.set()  # 设置事件，通知主线程
                break
    except asyncio.CancelledError:
        logging.info(f"Client {client_id} check task cancelled.")


async def check_client_online_signal(client_id, websocket, disconnect_event) -> bool:
    try:
        await websocket.ping()  # 发送ping消息
        # logging.info(f"Client {client_id} is still online.")
        return True
    except websockets.exceptions.ConnectionClosed:
        logging.info(f"Client {client_id} is offline.")
        disconnect_event.set()  # 设置事件，通知主线程
        return False


async def handle_client_connection(websocket):
    # 为每个客户端生成唯一标识符
    client_id = str(uuid.uuid4())
    connected_clients[client_id] = websocket
    logging.info(f"Client {client_id} connected.")

    # 存储客户端的事件\信息
    client_info[client_id] = {'disconnect_event': asyncio.Event(),  # 创建一个事件来通知连接断开
                              'require_file_event': asyncio.Event(),  # 创建一个事件来fetch文件
                              'require_file_tree_event': asyncio.Event(),  # 创建一个事件来fetch文件树
                              'require_file_tree_json_done_event': asyncio.Event(),  # 等待文件树建立响应完成
                              'require_file_pause_event': asyncio.Event(),  # 创建一个事件来暂停文件传输
                              'require_file_resume_event': asyncio.Event(),  # 创建一个事件来恢复文件传输
                              'filename': None,  # 存储客户端请求的文件名
                              'client_file_lists': set()  # 存储客户端的文件列表
                              }
    # 启动后台任务检查客户端是否在线
    check_client_online_task = asyncio.create_task(
        check_client_online(client_id, websocket, client_info[client_id]['disconnect_event']))

    try:
        while True:

            # 使用 asyncio.wait 同时等待多个任务：断开连接Event or 发送文件Event
            await asyncio.wait([
                client_info[client_id]['disconnect_event'].wait(),
                client_info[client_id]['require_file_event'].wait(),
                client_info[client_id]['require_file_tree_event'].wait(),
                client_info[client_id]['require_file_pause_event'].wait(),
                client_info[client_id]['require_file_resume_event'].wait()
            ], return_when=asyncio.FIRST_COMPLETED)
            # Situation 1. 检查是否收到断开连接的通知
            if client_info[client_id]['disconnect_event'].is_set():
                logging.info(f"Client {client_id} disconnected detected by ping task.")
                break

            # Situation 2. 检查是否收到发送文件的通知
            elif client_info[client_id]['require_file_event'].is_set():
                # 1. 发送请求文件的名字
                file_name = client_info[client_id]['filename']
                send_json = json.dumps({'type': 'request_file',
                                        'file_name': file_name})
                await websocket.send(send_json)
                # 等待客户端反馈：文件是否存在
                response = await websocket.recv()
                response_data = json.loads(response)

                # 如果要请求文件
                if response_data['type'] == 'file_info' and response_data['file_exists'] == True:
                    # 反馈本地文件的range.json给client -> 初始化本地 range.json
                    file_range_json = response_data['range_data']
                    server_range_json = {
                        "filename": file_range_json.get("filename", 0),  # 文件名
                        "supports_resume": file_range_json.get("supports_resume", False),
                        "file_size": file_range_json.get("file_size", 0),
                        "total_chunks": file_range_json.get("chunk_count", 0),
                        "downloaded_chunks": [],
                        "file_hash": file_range_json.get("file_hash", ""),  # 保存服务器返回的哈希值
                        "download_accomplished": False  # 下载完成标志
                    }
                    sql_data = msql.get_upload_link(file_range_json['file_hash'], search_by='file_hash')
                    if sql_data is not None:  # 本地维护的range.json文件存在，使用本地的
                        if not sql_data['url_using']:  # URL过期 生成新的URL
                            link_id, expires_time = basicFuncation.generate_upload_link(file_name,
                                                                                        secret_key='123456789')
                            msql.update_upload_link(file_range_json['file_hash'], link_id=link_id,
                                                    expires_at=expires_time)
                        else:  # URL未过期 使用本地的
                            link_id = sql_data['link_id']

                        if sql_data['file_chunk'] != file_range_json['chunk_count']:
                            # 服务器和客户端的文件块数量不一致
                            print(
                                f"服务器和客户端的文件块数量不一致{sql_data['file_chunk']}, {file_range_json['chunk_count']}")
                            sql_data['file_chunk'] = file_range_json['chunk_count']  # 更新服务器文件块数量
                            sql_data['downloaded_chunks'] = []  # 重传文件 清空已下载块
                            server_range_json['downloaded_chunks'] = []
                            msql.update_upload_link(file_range_json['file_hash'],
                                                    file_chunk=sql_data['file_chunk'],
                                                    downloaded_chunks=json.dumps([]))
                        else:
                            server_range_json['downloaded_chunks'] = json.loads(sql_data['downloaded_chunks'])
                        logging.info(f'link_id {link_id}')
                    else:  # 数据库维护的range.json文件不存在，立即创建
                        link_id, expires_time = basicFuncation.generate_upload_link(file_name, secret_key='123456789')
                        # 保存server_range_json 和 file_range_json 到 upload_links 中
                        msql.insert_upload_link(file_hash=file_range_json['file_hash'],
                                                file_name=file_range_json['filename'],
                                                file_size=file_range_json['file_size'],
                                                file_chunk=file_range_json['chunk_count'],
                                                link_id=link_id,
                                                expires_at=expires_time,
                                                url_using=True,
                                                supports_resume=file_range_json['supports_resume'],
                                                downloaded_chunks=json.dumps([]))
                        msql.insert_local_files(file_hash=file_range_json['file_hash'],
                                                file_name=file_range_json['filename'],
                                                file_exist=False,
                                                file_size=file_range_json['file_size'],
                                                file_path='',
                                                upload_time=datetime.now())

                    # 2. 创建空洞文件
                    if not os.path.exists(os.path.join(UPLOAD_FOLDER, file_range_json['file_hash'])):  # 文件存在则跳过
                        with open(os.path.join(UPLOAD_FOLDER, file_range_json['file_hash']), 'wb') as f:
                            f.seek(file_range_json['file_size'] - 1)
                            f.write(b'\x00')
                    # 3. 发送server_range.json和sendto_url给client
                    await websocket.send(json.dumps({'type': 'server_range_json',
                                                     'sendto_url': f'http://{SERVER_IP}:5000/upload/{link_id}',
                                                     'server_range_json': server_range_json}))

                    # 3. 文件开始发送 -> 清空文件请求事件 and 清空文件暂停传输事件
                    client_info[client_id]['require_file_event'].clear()
                    client_info[client_id]['require_file_pause_event'].clear()

            # Situation 3. 检查是否请求文件树
            elif client_info[client_id]['require_file_tree_event'].is_set():
                send_json = json.dumps({'type': 'request_file_tree'})
                await websocket.send(send_json)
                response = await websocket.recv()
                client_info[client_id]['client_file_lists'] = response  # 是一个json
                client_info[client_id]['require_file_tree_json_done_event'].set()  # 通知client_file_lists已经获取完成
                client_info[client_id]['require_file_tree_event'].clear()  # 清空文件树请求事件 等待下一次请求

            # Situation 4. 检查是否停止文件传输
            elif client_info[client_id]['require_file_pause_event'].is_set():
                logger.info("暂停文件上传")
                send_json = json.dumps({'type': 'require_file_pause'})
                await websocket.send(send_json)
                client_info[client_id]['require_file_pause_event'].clear()

            # Situation 4. 检查是否停止文件传输
            elif client_info[client_id]['require_file_resume_event'].is_set():
                logger.info("恢复文件上传")
                send_json = json.dumps({'type': 'require_file_resume'})
                await websocket.send(send_json)
                client_info[client_id]['require_file_resume_event'].clear()

            else:
                logging.info(f"Event not found for client {client_id}")

    except websockets.exceptions.ConnectionClosed:
        logging.info(f"Client {client_id} disconnected.")
    finally:
        # 清理断开连接的客户端
        if client_id in connected_clients:
            del connected_clients[client_id]
            logging.info(f"Client {client_id} removed from connected_clients.")
        # 取消后台任务
        check_client_online_task.cancel()


async def handle_monitor_connection(websocket):
    """处理监控连接"""
    while True:
        try:
            # 接收消息
            message = await websocket.recv()
            data = json.loads(message)
            # logger.info(data)
            if data['type'] == 'get_clients_info':
                # 处理客户端信息请求
                logging.info("处理客户端信息请求")
                clients_list = []
                # 判断客户端是否在线 and 构建clients_list
                for client_id in connected_clients.keys():
                    if await check_client_online_signal(client_id,
                                                        connected_clients[client_id],
                                                        client_info[client_id]['disconnect_event']):
                        clients_list.append({"client_id": client_id})
                await websocket.send(json.dumps({"type": "clients_info", "clients": clients_list}))
            elif data['type'] == 'get_file_tree':
                # 处理文件树请求
                logging.info(f"Client {data['client_id']} 请求文件树")
                client_id = data['client_id']
                if (client_id in connected_clients.keys() and
                        not await check_client_online_signal(client_id,
                                                             connected_clients[client_id],
                                                             client_info[client_id]['disconnect_event'])):
                    continue  # 如果客户端不在线了
                # 获取文件树json
                client_info[client_id]['require_file_tree_json_done_event'].clear()
                client_info[client_id]['require_file_tree_event'].set()
                await client_info[client_id]['require_file_tree_json_done_event'].wait()
                directory_structure = json.loads(client_info[client_id]['client_file_lists'])

                # 检查每个文件是否存在
                def add_exists_flag(node):
                    if node['type'] == 'file':
                        file_name = node['name']
                        local_file_info = msql.get_local_files(file_name, search_by='file_name')
                        if local_file_info is None:
                            node['exists'] = False
                        else:
                            file_path = os.path.join(SERVER_FILES_DIR, local_file_info['file_hash'])
                            node['exists'] = os.path.exists(file_path)
                    elif node['type'] == 'directory' and node['children']:
                        for child in node['children']:
                            add_exists_flag(child)

                add_exists_flag(directory_structure)
                await websocket.send(json.dumps({
                    "type": "file_tree",
                    "client_id": client_id,
                    "files": directory_structure
                }))
            elif data['type'] == 'file_click':
                # 处理文件点击
                client_id = data['client_id']
                if (client_id in connected_clients.keys() and
                        not await check_client_online_signal(client_id,
                                                             connected_clients[client_id],
                                                             client_info[client_id]['disconnect_event'])):
                    continue  # 如果客户端不在线了
                client_info[data['client_id']]['require_file_event'].set()
                client_info[data['client_id']]['filename'] = data['filename']
                logging.info(f"Client {data['client_id']} 点击了文件: {data['filename']}")
            elif data['type'] == 'file_upload_progress':
                client_id = data['client_id']
                file_path = data['file_path']
                file_name = os.path.basename(file_path)
                # logger.info(f"Client {client_id} 获取上传进度: {file_name}")
                while msql.get_local_files(file_name, search_by='file_name') is None:
                    await asyncio.sleep(0.02)
                file_hash = msql.get_local_files(file_name, search_by='file_name')['file_hash']
                # 从数据库获取上传进度
                sql_data = msql.get_upload_link(file_hash, search_by='file_hash')
                if sql_data is not None:  # 如果有记录
                    total_chunks = sql_data['file_chunk']
                    downloaded_chunks = len(json.loads(sql_data['downloaded_chunks']))
                    progress = downloaded_chunks / total_chunks * 100
                    await websocket.send(json.dumps({
                        "type": "upload_progress",
                        "client_id": client_id,
                        "file_name": file_name,
                        "file_hash": file_hash,
                        "progress": progress
                    }))
                else:
                    await websocket.send(json.dumps({
                        "type": "upload_error",
                        "message": "File not found"
                    }))
            elif data['type'] == 'get_server_files':
                # 获取服务器文件列表
                files = []
                for root, dirs, files_in_dir in os.walk(SERVER_FILES_DIR):
                    for file_hash in files_in_dir:
                        file_path = os.path.join(root, file_hash)
                        file_stat = os.stat(file_path)
                        local_file_info = msql.get_local_files(file_hash, search_by='file_hash')
                        if local_file_info is not None:  # 如果文件在数据库中存在
                            file_name = local_file_info['file_name']
                        else:
                            continue
                        files.append({
                            'hash': file_hash,
                            'name': file_name,
                            'path': file_path,
                            'size': basicFuncation.format_size(file_stat.st_size),
                            'modified': datetime.fromtimestamp(file_stat.st_mtime).strftime('%Y-%m-%d %H:%M:%S')
                        })
                await websocket.send(json.dumps({
                    "type": "server_files",
                    "files": files
                }))
            elif data['type'] == 'pause_upload':
                # 处理暂停上传请求
                file_hash = data['file_hash']
                sql_data = msql.get_upload_link(file_hash, search_by='file_hash')
                if sql_data:
                    # 更新数据库中的上传状态
                    logger.info(f"暂停上传文件: {file_hash}")
                    client_info[data['client_id']]['require_file_pause_event'].set()
                    await websocket.send(json.dumps({
                        "type": "pause_upload_result",
                        "success": True,
                        "message": "上传已暂停"
                    }))
                else:
                    logger.info(f"文件未找到: {file_hash}")
                    client_info[data['client_id']]['require_file_resume_event'].set()
                    await websocket.send(json.dumps({
                        "type": "pause_upload_result",
                        "success": False,
                        "message": "文件未找到"
                    }))
            elif data['type'] == 'resume_upload':
                # 处理恢复上传请求
                file_hash = data['file_hash']
                sql_data = msql.get_upload_link(file_hash, search_by='file_hash')
                if sql_data:
                    # 更新数据库中的上传状态
                    logger.info(f"恢复上传文件: {file_hash}")
                    client_info[data['client_id']]['require_file_resume_event'].set()
                    await websocket.send(json.dumps({
                        "type": "resume_upload_result",
                        "success": True,
                        "message": "上传已恢复"
                    }))
                else:
                    logger.info(f"文件未找到: {file_hash}")
                    await websocket.send(json.dumps({
                        "type": "resume_upload_result",
                        "success": False,
                        "message": "文件未找到"
                    }))
            elif data['type'] == 'delete_file':
                # 删除服务器文件
                file_path = data['file_path']
                try:
                    if os.path.exists(file_path):
                        os.remove(file_path)
                        # 从数据库中删除文件记录
                        file_hash = os.path.basename(file_path)
                        msql.update_local_files(file_hash, file_exist=False)
                        await websocket.send(json.dumps({
                            "type": "delete_file_result",
                            "success": True,
                            "message": "文件删除成功"
                        }))
                    else:
                        await websocket.send(json.dumps({
                            "type": "delete_file_result",
                            "success": False,
                            "message": "文件不存在"
                        }))
                except Exception as e:
                    await websocket.send(json.dumps({
                        "type": "delete_file_result",
                        "success": False,
                        "message": f"文件删除失败: {str(e)}"
                    }))
            else:
                logger.info(f"未知的消息类型: {data['type']}")
        except websockets.exceptions.ConnectionClosed:
            logging.info("Monitor WebSocket connection closed.")
            break


async def monitor_local_files():
    """
    监控本地文件变化
    :return:
    """
    while True:
        try:
            msql.check_FilesExist(SERVER_FILES_DIR)
            msql.check_UrlAvailable()
        except Exception as e:
            logging.info(f"File monitoring error: {e}")
        finally:
            pass
        await asyncio.sleep(2)


async def run_websocket():
    # 启动文件监控任务
    monitor_task = asyncio.create_task(monitor_local_files())

    # 启动主WebSocket服务器
    async with websockets.serve(handle_client_connection, SERVER_IP, 8765):
        logging.info(f"WebSocket server started on ws://{SERVER_IP}:8765")

        # 启动监控WebSocket服务器
        async with websockets.serve(handle_monitor_connection, SERVER_IP, 8766):
            logging.info(f"Monitor WebSocket server started on ws://{SERVER_IP}:8766")
            await asyncio.Future()  # 保持服务器运行


if __name__ == "__main__":
    db = msql.Database()
    # 确保服务器文件目录存在
    os.makedirs(SERVER_FILES_DIR, exist_ok=True)

    # 配置Windows事件循环
    # if os.name == 'nt':
    #     asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

    # 启动Websocket服务器
    asyncio.run(run_websocket())
